import pyfpgrowth
# import fp_growth
import pygraphviz as pgv
from collections import defaultdict
import csv
# 读取数据集
def load_dataset(file_path):
    dataset = []
    # with open(file_path, 'r') as file:
    #     for line in file:
    #         dataset.append([item.strip() for item in line.split(',')])
    with open(file_path, 'r') as file:
        reader = csv.DictReader(file)
        for row in reader:
            # 跳过空单元格
            items = [item.strip() for item in row.values() if item.strip()]
            if items:
                dataset.append(items)
    print(dataset)
    return dataset

# 构建FP树
def build_fp_tree(dataset, min_support):
    patterns = pyfpgrowth.find_frequent_patterns(dataset, min_support)
    sorted_patterns = {k: v for k, v in sorted(patterns.items(), key=lambda item: item[1], reverse=True)}
    tree = defaultdict(dict)
    for transaction, count in sorted_patterns.items():
        current = tree
        for item in transaction:
            if item not in current:
                current[item] = {'count': count, 'children': {}}
            else:
                current[item]['count'] += count
            current = current[item]['children']
    return tree

# 可视化FP树
def visualize_fp_tree(tree,rules, output_file='fp_tree_visualization.png'):
    G = pgv.AGraph(strict=False, directed=True)
    G.graph_attr['rankdir'] = 'TB'
    G.add_node('null', label='', style='invisible')
    build_graph_with_rules(G, 'null', tree,rules)
    G.layout(prog='dot')
    G.draw(output_file,format='svg')

def build_graph(G, parent, tree):
    for item, values in tree.items():
        node_id = f"{parent}_{item}"
        G.add_node(node_id, label=f"{item} ({values['count']})")
        G.add_edge(parent, node_id)
        build_graph(G, node_id, values['children'])

def build_graph_with_rules(G, parent, tree,rules):
    """构建图形节点和边的函数"""
    for item, values in tree.items():
        node_id = f"{parent}_{item}"
        G.add_node(node_id, label=f"{item} ({values['count']})")
        for rule_itemset,(right_itemset,confidence)in rules.items():

            if set(rule_itemset).issubset(set(node_id.split('_')[1:])):
                right_item=right_itemset[0] if len(right_itemset)==1 else tuple(right_itemset)
                right_node_id=f"{node_id}_{right_item}"
                G.add_node(right_node_id,label=f"({values['count']})",style='filled',fillcolor='lightblue')
                G.add_edge(node_id,right_node_id,label=f"conf={confidence:.2f}")
        G.add_edge(parent, node_id)
        build_graph_with_rules(G, node_id, values['children'],rules)


# 加载数据集
dataset = load_dataset("D:\Lenovo\Desktop\云南大学\空间数据挖掘\实验数据\实验数据5.csv")

# 构建FP树
min_support = 100 #最小支持度计数，即出现某个项集的频次
confidence=0.8 #置信度阈值
fp_tree = build_fp_tree(dataset, min_support)


# 生成关联规则
patterns = pyfpgrowth.find_frequent_patterns(dataset, min_support)
rules=pyfpgrowth.generate_association_rules(patterns, confidence)
print(rules)
# 可视化FP树
visualize_fp_tree(fp_tree, rules,output_file='fp_tree_visualization.svg')
print('done')

